
import sys
import indexing.featureindex

sys.path.append("/archive10/bowtie/python_scripts")
import sampleParams

def readPhastConsValues(file="/home/bowtie/Bowtie_Illumina_Alignment/references/phastCons/breast-pair-snp-locations.phastcons.txt"):
    phastcons = {}
    for line in open(file):
        cols = line.strip().split("\t")
        phastcons[cols[0]] = cols[1]
    print >> sys.stderr, "# Read %d entries from PhastCons file" % len(phastcons)
    return phastcons

def combineAcrossClasses(file, excludedSamples=[]):
    phastcons = readPhastConsValues()
    params = sampleParams.getCurrentSampleParams()
    sampleClass = {}
    breastNormal = []
    breastCancer = []
    nonBreastNormal = []
    nonBreastCancer = []
    for paramName in params.keys():
        p = params[paramName]
        sample = p.getParam("Sample")
        tissue = p.getProperty("TissueType", "NA")
        treatment = p.getProperty("Treatment", "NA")
        progression = p.getProperty("Progression", "NA")
        if treatment != "FALSE": continue
        if tissue == "NA" or (progression == "NA" and not tissue == "B Lymphocyte"): continue
        isBreast = (tissue == "Breast" or tissue == "B Lymphocyte")
        isCancer = (not progression.startswith("Benign") and not tissue == "B Lymphocyte")
        if isBreast:
            if isCancer:
                sampleClass[sample] = "Breast Cancer"
                breastCancer.append(sample)
            else:
                sampleClass[sample] = "Breast Normal"
                breastNormal.append(sample)
        else:
            if isCancer:
                sampleClass[sample] = "Non-Breast Cancer"
                nonBreastCancer.append(sample)
            else:
                sampleClass[sample] = "Non-Breast Normal"
                nonBreastNormal.append(sample)
    classSizes = {}
    print >> sys.stdout, "# %d breast normal: " % len(breastNormal), breastNormal
    classSizes["Breast Normal"] = len(breastNormal)
    print >> sys.stdout, "# %d breast cancer: " % len(breastCancer), breastCancer
    classSizes["Breast Cancer"] = len(breastCancer)
    print >> sys.stdout, "# %d non-breast normal: " % len(nonBreastNormal), nonBreastNormal
    classSizes["Non-Breast Normal"] = len(nonBreastNormal)
    print >> sys.stdout, "# %d non-breast cancer: " % len(nonBreastCancer), nonBreastCancer
    classSizes["Non-Breast Cancer"] = len(nonBreastCancer)
    
    sampleClasses = ["Breast Normal", "Breast Cancer", "Non-Breast Normal", "Non-Breast Cancer"]

    for line in open(file):
        cols = line.strip().split("\t")
        if len(cols)==1:
            samplecounts = []
        else:
            samplecounts = cols[1].split(",")
        totalCoverage = {}
        totalBaseCounts = {}
        totalSampleAnyCounts = {}
        totalSampleAnyExpression = {}
        #totalSampleGoodCounts = {}
        #totalSampleGoodExpression = {}

        maxReads = {}
        maxPercent = {}

        for cls in sampleClasses:
            totalBaseCounts[cls] = [0,0,0,0]
            totalCoverage[cls] = 0
            totalSampleAnyCounts[cls] = [0,0,0,0]
            totalSampleAnyExpression[cls] = 0
            #totalSampleGoodCounts[cls] = [0,0,0,0]
            #totalSampleGoodExpression[cls] = 0
            maxReads[cls] = [("NA", 0, 0) for i in range(4)]
            maxPercent[cls] = [("NA", 0.0, 0, 0) for i in range(4)]

        for ct in samplecounts:
            (sample, basecounts) = tuple(ct.split("="))
            if not sampleClass.has_key(sample): continue
            cls = sampleClass[sample]
            bases = basecounts.split("/")
            cov = sum([int(bases[ix]) for ix in xrange(4)])
            if cov>=1:
                totalSampleAnyExpression[cls] += 1
            if cov>=4:
                pass # totalSampleGoodExpression[cls] += 1
            totalCoverage[cls] += cov
            for ix in xrange(4):
                if int(bases[ix])>0:
                    totalBaseCounts[cls][ix] += int(bases[ix])
                    totalSampleAnyCounts[cls][ix] += 1
                    if cov >= 4 and 1.0*int(bases[ix])/cov >= 0.3:
                        pass # totalSampleGoodCounts[cls][ix] += 1
                    if int(bases[ix]) > maxReads[cls][ix][1] and not sample in excludedSamples:
                        maxReads[cls][ix] = (sample, int(bases[ix]), cov)
                    if cov >= 4 and 100.0*int(bases[ix])/cov > maxPercent[cls][ix][1] and not sample in excludedSamples:
                        maxPercent[cls][ix] = (sample, 100.0*int(bases[ix])/cov, int(bases[ix]), cov)

        bases = ["A", "C", "G", "T"]

        for ix in xrange(4):
            baseCountsStr = "\t".join([str(totalBaseCounts[cls][ix]) for cls in sampleClasses])
            coverageStr = "\t".join([str(totalCoverage[cls]) for cls in sampleClasses])
            sampleAnyCountsStr = "\t".join([str(totalSampleAnyCounts[cls][ix]) for cls in sampleClasses])
            sampleAnyExpressionStr = "\t".join([str(totalSampleAnyExpression[cls]) for cls in sampleClasses])
            classSizeStr = "\t".join([str(classSizes[cls]) for cls in sampleClasses])
            #sampleGoodCountsStr = "\t".join([str(totalSampleGoodCounts[cls][ix]) for cls in sampleClasses])
            #sampleGoodExpressionStr = "\t".join([str(totalSampleGoodExpression[cls]) for cls in sampleClasses])
            maxReadsStr = "\t".join(["%s\t%d\t%d" % maxReads[cls][ix] for cls in sampleClasses])
            maxPercentStr = "\t".join(["%s\t%.2f\t%d\t%d" % maxPercent[cls][ix] for cls in sampleClasses])
            
            # st = "%s\t%s\t%s\t%s\t%s\t%s" % (baseCountsStr, coverageStr, sampleAnyCountsStr, sampleAnyExpressionStr, sampleGoodCountsStr, sampleGoodExpressionStr)
            st = "%s\t%s\t%s\t%s\t%s\t%s\t%s" % (baseCountsStr, coverageStr, sampleAnyCountsStr, sampleAnyExpressionStr, classSizeStr, maxReadsStr, maxPercentStr)
            print "%s\t%s\t%s\t%s" % (cols[0], phastcons.get(cols[0], "Unknown"), bases[ix], st)

def combineAllCounts(files):
    allLocs = {}
    locs = {}
    print >> sys.stderr, "# Processing %d files" % len(files)
    for file in files:
        print >> sys.stderr, "# Processing file # %d: %s" % (files.index(file), file)
        allLocs = {}
        for line in open(file):
            cols = line.strip().split("\t")
            allLocs[cols[0]] = True
            if len(cols)==1:
                continue
            if locs.has_key(cols[0]):
                locs[cols[0]].append(cols[1])
            else:
                locs[cols[0]] = [cols[1]]
        print >> sys.stderr, "# %d positions" % len(allLocs)
    for loc in allLocs.keys():
        if locs.has_key(loc):
            print "%s\t%s" % (loc, ",".join(locs[loc]))
        else:
            print loc
    
def addSampleCounts(locfile, sample):
    locs = {}
    for line in open(locfile):
        cols = line.strip().split("\t")
        (chr, pos) = cols[0].split(":")
        pos = long(pos)
        locs[(chr, pos)] = ""
    print >> sys.stderr, "# Read %ld locs from %s" % (len(locs), locfile)
    if len(locs)<1000:
        coverageIndex = indexing.featureindex.FeatureIndex("/archive10/bowtie/Sample_Parameters/%s/Current/%s.coverage_summary.INDEX" % (sample, sample))
        numFeatures = coverageIndex.getFeatureCount()
        numRead = 0
        for (chr, pos) in locs.keys():
            for feature in coverageIndex.getFeatures((chr, pos, pos)):
                numRead += 1
                if int(100.0*numRead/numFeatures)>int(100.0*(numRead-1)/numFeatures):
                    print >> sys.stderr, "# Processed %d of %d total features; %d locations remaining" % (numRead, numFeatures, len(locs))
                key = (feature.chr, feature.spos)
                if not locs.has_key(key): continue
                details = feature.getDetails()
                if locs[key]!="": locs[key]+=","
                locs[key] += "%s=%s" % (sample, details.getCountsAsString())
                print "%s:%ld\t%s" % (key[0], key[1], locs[key])
                del locs[key]
        print >> sys.stderr, "# %ld locations remaining" % len(locs)
        for loc in locs.keys():
            print "%s:%ld\t%s" % (loc[0], loc[1], locs[loc])
        return
    coverageIndex = indexing.featureindex.FeatureIndex("/archive10/bowtie/Sample_Parameters/%s/Current/%s.coverage_summary.INDEX" % (sample, sample))
    numFeatures = coverageIndex.getFeatureCount()
    numRead = 0
    for feature in coverageIndex.getFeatures():
        numRead += 1
        if int(100.0*numRead/numFeatures)>int(100.0*(numRead-1)/numFeatures):
            print >> sys.stderr, "# Processed %d of %d total features; %d locations remaining" % (numRead, numFeatures, len(locs))
        key = (feature.chr, feature.spos)
        if not locs.has_key(key): continue
        details = feature.getDetails()
        if locs[key]!="": locs[key]+=","
        locs[key] += "%s=%s" % (sample, details.getCountsAsString())
        print "%s:%ld\t%s" % (key[0], key[1], locs[key])
        del locs[key]
    print >> sys.stderr, "# %ld locations remaining" % len(locs)
    for loc in locs.keys():
        print "%s:%ld\t%s" % (loc[0], loc[1], locs[loc])

if __name__=="__main__":
    if sys.argv[1]=="-S":
        locfile = sys.argv[2]
        sample = sys.argv[3]
        addSampleCounts(locfile, sample)
    elif sys.argv[1]=="-A":
        combineAcrossClasses(sys.argv[2], sys.argv[3:])
    else:
        combineAllCounts(sys.argv[1:])

